import numpy as np
import statsmodels.formula.api as smf
from sklearn.model_selection import train_test_split
from sqlalchemy import create_engine
import pymysql
from sklearn.metrics import r2_score, mean_absolute_error
import pandas as pd

# 数据库配置
db_conf = {
    'host': '111.231.14.211',
    'user': 'tushare',
    'password': 'root',
    'database': 'tushare',
    'port': 13307,
    'charset': 'utf8mb4',
}

# 创建引擎与连接
engine = create_engine(
    f"mysql+pymysql://{db_conf['user']}:{db_conf['password']}@{db_conf['host']}:{db_conf['port']}/{db_conf['database']}")
conn = pymysql.connect(**db_conf)

# 查询数据
query = """
    SELECT d.* FROM date_1 d 
    WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' 
      AND d.ts_code = '000001.SZ'
"""

df1 = pd.read_sql_query(query, conn)
conn.close()

# 计算涨跌幅
df1['zd_closes'] = df1['closes'].pct_change().round(4)  # 使用 pct_change 更简洁

# 处理缺失值
df1 = df1.dropna(subset=['zd_closes'])

# 特征选择
exclude_cols = ['id', 'ts_code', 'trade_date', 'the_date',
                'opens', 'high', 'low', 'closes', 'pre_closes', 'changes', 'pct_chg', 'zd_closes']

number_cols = df1.select_dtypes(include=['number']).columns.tolist()
features = [col for col in number_cols if col not in exclude_cols]

# 构建公式
formula = 'zd_closes ~ ' + ' + '.join(features)

# 划分训练集和测试集
df_train, df_test = train_test_split(df1, test_size=0.2, random_state=42)

# 建立回归模型
model = smf.ols(formula, data=df_train).fit()

# 预测与评估
preds = model.predict(df_test)
r2 = r2_score(df_test['zd_closes'], preds)
mae = mean_absolute_error(df_test['zd_closes'], preds)

# 输出结果
print(model.summary())
print(f"\nTest Set Performance:")
print(f"R² Score: {r2:.4f}, MAE: {mae:.6f}")